{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "14ec0fd7",
   "metadata": {},
   "source": [
    "# Finetuning a DistilBERT Classifier in Lightning with `finetuning-scheduler`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56f0ce0a",
   "metadata": {},
   "source": [
    "![](figures/finetuning-ii.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c70f0d1-fae7-4f33-a119-63ab47739f27",
   "metadata": {},
   "source": [
    "Here, we are using the [finetuning-scheduler](https://github.com/speediedan/finetuning-scheduler) callback for freezing and unfreezing layers during finetuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7133d6ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "32c47e21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5832dbd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install pytorch_lightning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "58c34189",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install finetuning-scheduler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8f06f3c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-11-16 19:25:15.509473: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch               : 1.12.0\n",
      "transformers        : 4.24.0\n",
      "datasets            : 2.7.0\n",
      "pytorch_lightning   : 1.8.1\n",
      "finetuning_scheduler: 0.3.1\n",
      "\n",
      "conda environment: base\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%load_ext watermark\n",
    "%watermark --conda -p torch,transformers,datasets,pytorch_lightning,finetuning_scheduler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4cab474",
   "metadata": {},
   "source": [
    "# 1 Loading the Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f099ed3",
   "metadata": {},
   "source": [
    "The IMDB movie review dataset consists of 50k movie reviews with sentiment label (0: negative, 1: positive)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7070196e",
   "metadata": {},
   "source": [
    "## 1a) Load from `datasets` Hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c58a5bee",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import list_datasets, load_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ba04345f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# list_datasets()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f99da96b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset imdb (/home/jovyan/.cache/huggingface/datasets/imdb/plain_text/1.0.0/2fdd8b9bcadd6e7055e742a706876ba43f19faee861df134affd7a3f60fc38a1)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "180fa99b0aa24d84942f8cd6d0165a2e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['text', 'label'],\n",
      "        num_rows: 25000\n",
      "    })\n",
      "    test: Dataset({\n",
      "        features: ['text', 'label'],\n",
      "        num_rows: 25000\n",
      "    })\n",
      "    unsupervised: Dataset({\n",
      "        features: ['text', 'label'],\n",
      "        num_rows: 50000\n",
      "    })\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "imdb_data = load_dataset(\"imdb\")\n",
    "print(imdb_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7c80ddb7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'text': \"This film is terrible. You don't really need to read this review further. If you are planning on watching it, suffice to say - don't (unless you are studying how not to make a good movie).<br /><br />The acting is horrendous... serious amateur hour. Throughout the movie I thought that it was interesting that they found someone who speaks and looks like Michael Madsen, only to find out that it is actually him! A new low even for him!!<br /><br />The plot is terrible. People who claim that it is original or good have probably never seen a decent movie before. Even by the standard of Hollywood action flicks, this is a terrible movie.<br /><br />Don't watch it!!! Go for a jog instead - at least you won't feel like killing yourself.\",\n",
       " 'label': 0}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imdb_data[\"train\"][99]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90af222a",
   "metadata": {},
   "source": [
    "## 1b) Load from local directory"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51b3b63d",
   "metadata": {},
   "source": [
    "The IMDB movie review set can be downloaded from http://ai.stanford.edu/~amaas/data/sentiment/. After downloading the dataset, decompress the files.\n",
    "\n",
    "A) If you are working with Linux or MacOS X, open a new terminal window cd into the download directory and execute\n",
    "\n",
    "    tar -zxf aclImdb_v1.tar.gz\n",
    "\n",
    "B) If you are working with Windows, download an archiver such as 7Zip to extract the files from the download archive."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c409d5e",
   "metadata": {},
   "source": [
    "C) Use the following code to download and unzip the dataset via Python"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7220fff4",
   "metadata": {},
   "source": [
    "**Download the movie reviews**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c9fd3127",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import tarfile\n",
    "import time\n",
    "import urllib.request\n",
    "\n",
    "source = \"http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz\"\n",
    "target = \"aclImdb_v1.tar.gz\"\n",
    "\n",
    "if os.path.exists(target):\n",
    "    os.remove(target)\n",
    "\n",
    "\n",
    "def reporthook(count, block_size, total_size):\n",
    "    global start_time\n",
    "    if count == 0:\n",
    "        start_time = time.time()\n",
    "        return\n",
    "    duration = time.time() - start_time\n",
    "    progress_size = int(count * block_size)\n",
    "    speed = progress_size / (1024.0**2 * duration)\n",
    "    percent = count * block_size * 100.0 / total_size\n",
    "\n",
    "    sys.stdout.write(\n",
    "        f\"\\r{int(percent)}% | {progress_size / (1024.**2):.2f} MB \"\n",
    "        f\"| {speed:.2f} MB/s | {duration:.2f} sec elapsed\"\n",
    "    )\n",
    "    sys.stdout.flush()\n",
    "\n",
    "\n",
    "if not os.path.isdir(\"aclImdb\") and not os.path.isfile(\"aclImdb_v1.tar.gz\"):\n",
    "    urllib.request.urlretrieve(source, target, reporthook)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8a140908",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.isdir(\"aclImdb\"):\n",
    "\n",
    "    with tarfile.open(target, \"r:gz\") as tar:\n",
    "        tar.extractall()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20917799",
   "metadata": {},
   "source": [
    "**Convert them to a pandas DataFrame and save them as CSV**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6491b911",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50000/50000 [01:21<00:00, 614.13it/s]\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from packaging import version\n",
    "from tqdm import tqdm\n",
    "\n",
    "# change the `basepath` to the directory of the\n",
    "# unzipped movie dataset\n",
    "\n",
    "basepath = \"aclImdb\"\n",
    "\n",
    "labels = {\"pos\": 1, \"neg\": 0}\n",
    "\n",
    "df = pd.DataFrame()\n",
    "\n",
    "with tqdm(total=50000) as pbar:\n",
    "    for s in (\"test\", \"train\"):\n",
    "        for l in (\"pos\", \"neg\"):\n",
    "            path = os.path.join(basepath, s, l)\n",
    "            for file in sorted(os.listdir(path)):\n",
    "                with open(os.path.join(path, file), \"r\", encoding=\"utf-8\") as infile:\n",
    "                    txt = infile.read()\n",
    "\n",
    "                if version.parse(pd.__version__) >= version.parse(\"1.3.2\"):\n",
    "                    x = pd.DataFrame(\n",
    "                        [[txt, labels[l]]], columns=[\"review\", \"sentiment\"]\n",
    "                    )\n",
    "                    df = pd.concat([df, x], ignore_index=False)\n",
    "\n",
    "                else:\n",
    "                    df = df.append([[txt, labels[l]]], ignore_index=True)\n",
    "                pbar.update()\n",
    "df.columns = [\"text\", \"label\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "890e08eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "np.random.seed(0)\n",
    "df = df.reindex(np.random.permutation(df.index))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "949b9362",
   "metadata": {},
   "source": [
    "**Basic datasets analysis and sanity checks**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "7dcf9d12",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Class distribution:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([25000, 25000])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Class distribution:\")\n",
    "np.bincount(df[\"label\"].values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "537aea79",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4, 173.0, 2470)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text_len = df[\"text\"].apply(lambda x: len(x.split()))\n",
    "text_len.min(), text_len.median(), text_len.max() "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "934f4554",
   "metadata": {},
   "source": [
    "**Split data into training, validation, and test sets**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "4d2fa677",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_shuffled = df.sample(frac=1, random_state=1).reset_index()\n",
    "\n",
    "df_train = df_shuffled.iloc[:35_000]\n",
    "df_val = df_shuffled.iloc[35_000:40_000]\n",
    "df_test = df_shuffled.iloc[40_000:]\n",
    "\n",
    "df_train.to_csv(\"train.csv\", index=False, encoding=\"utf-8\")\n",
    "df_val.to_csv(\"validation.csv\", index=False, encoding=\"utf-8\")\n",
    "df_test.to_csv(\"test.csv\", index=False, encoding=\"utf-8\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e87ad72",
   "metadata": {},
   "source": [
    "# 2 Tokenization and Numericalization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d448a60b",
   "metadata": {},
   "source": [
    "**Load the dataset via `load_dataset`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "122ea695",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration default-d2824d6740ca7436\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading and preparing dataset csv/default to /home/jovyan/.cache/huggingface/datasets/csv/default-d2824d6740ca7436/0.0.0/6b34fb8fcf56f7c8ba51dc895bfa2bfbe43546f190a60fcf74bb5e8afdcc2317...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "02424be626814a81bff3c02499c3c8af",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data files:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   "
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "62c435a3f0bc474db33c1f9e10147bec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extracting data files #2:   0%|          | 0/1 [00:00<?, ?obj/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "238d65110d484f1bb9c09b6a643d511d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extracting data files #0:   0%|          | 0/1 [00:00<?, ?obj/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b62f93c0ebcd4b3e8aa1d7c38e66fbc6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extracting data files #1:   0%|          | 0/1 [00:00<?, ?obj/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating validation split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating test split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset csv downloaded and prepared to /home/jovyan/.cache/huggingface/datasets/csv/default-d2824d6740ca7436/0.0.0/6b34fb8fcf56f7c8ba51dc895bfa2bfbe43546f190a60fcf74bb5e8afdcc2317. Subsequent calls will reuse this data.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d4761a06ace04f6a98a4506c8493e543",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['index', 'text', 'label'],\n",
      "        num_rows: 35000\n",
      "    })\n",
      "    validation: Dataset({\n",
      "        features: ['index', 'text', 'label'],\n",
      "        num_rows: 5000\n",
      "    })\n",
      "    test: Dataset({\n",
      "        features: ['index', 'text', 'label'],\n",
      "        num_rows: 10000\n",
      "    })\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "imdb_dataset = load_dataset(\n",
    "    \"csv\",\n",
    "    data_files={\n",
    "        \"train\": \"train.csv\",\n",
    "        \"validation\": \"validation.csv\",\n",
    "        \"test\": \"test.csv\",\n",
    "    },\n",
    ")\n",
    "\n",
    "print(imdb_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d39a5ea8",
   "metadata": {},
   "source": [
    "**Tokenize the dataset**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d30070ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokenizer input max length: 512\n",
      "Tokenizer vocabulary size: 30522\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"distilbert-base-uncased\")\n",
    "print(\"Tokenizer input max length:\", tokenizer.model_max_length)\n",
    "print(\"Tokenizer vocabulary size:\", tokenizer.vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "a9a96a04",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_text(batch):\n",
    "    return tokenizer(batch[\"text\"], truncation=True, padding=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "32792e13",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cb7b69bebfb34ac7b4236c66f0b0ba5b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a623f934deaf495a851dca62796e728b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c492a8c311934347aa09e0f6a90fe1ba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "imdb_tokenized = imdb_dataset.map(tokenize_text, batched=True, batch_size=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "26b53ecc",
   "metadata": {},
   "outputs": [],
   "source": [
    "del imdb_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ef74024d",
   "metadata": {},
   "outputs": [],
   "source": [
    "imdb_tokenized.set_format(\"torch\", columns=[\"input_ids\", \"attention_mask\", \"label\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "1952d15f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7679c6ea",
   "metadata": {},
   "source": [
    "# 3 Set Up DataLoaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "22d94c8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "\n",
    "class IMDBDataset(Dataset):\n",
    "    def __init__(self, dataset_dict, partition_key=\"train\"):\n",
    "        self.partition = dataset_dict[partition_key]\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        return self.partition[index]\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.partition.num_rows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "40bbe57a",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = IMDBDataset(imdb_tokenized, partition_key=\"train\")\n",
    "val_dataset = IMDBDataset(imdb_tokenized, partition_key=\"validation\")\n",
    "test_dataset = IMDBDataset(imdb_tokenized, partition_key=\"test\")\n",
    "\n",
    "train_loader = DataLoader(\n",
    "    dataset=train_dataset,\n",
    "    batch_size=12,\n",
    "    shuffle=True, \n",
    "    num_workers=4\n",
    ")\n",
    "\n",
    "val_loader = DataLoader(\n",
    "    dataset=val_dataset,\n",
    "    batch_size=12,\n",
    "    num_workers=4\n",
    ")\n",
    "\n",
    "test_loader = DataLoader(\n",
    "    dataset=test_dataset,\n",
    "    batch_size=12,\n",
    "    num_workers=4\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d21316a",
   "metadata": {},
   "source": [
    "# 4 Initializing DistilBERT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "de14af0c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at distilbert-base-uncased were not used when initializing DistilBertForSequenceClassification: ['vocab_projector.bias', 'vocab_transform.weight', 'vocab_layer_norm.bias', 'vocab_projector.weight', 'vocab_layer_norm.weight', 'vocab_transform.bias']\n",
      "- This IS expected if you are initializing DistilBertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing DistilBertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of DistilBertForSequenceClassification were not initialized from the model checkpoint at distilbert-base-uncased and are newly initialized: ['pre_classifier.weight', 'classifier.bias', 'classifier.weight', 'pre_classifier.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModelForSequenceClassification\n",
    "\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    \"distilbert-base-uncased\", num_labels=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3409fb84",
   "metadata": {},
   "source": [
    "## 5 Finetuning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82bf5e5c",
   "metadata": {},
   "source": [
    "**Wrap in LightningModule for Training**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "b4f7f003",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytorch_lightning as L\n",
    "import torch\n",
    "import torchmetrics\n",
    "\n",
    "\n",
    "class LightningModel(L.LightningModule):\n",
    "    def __init__(self, model, learning_rate=5e-5):\n",
    "        super().__init__()\n",
    "\n",
    "        self.learning_rate = learning_rate\n",
    "        self.model = model\n",
    "\n",
    "        self.val_acc = torchmetrics.Accuracy()\n",
    "        self.test_acc = torchmetrics.Accuracy()\n",
    "\n",
    "    def forward(self, input_ids, attention_mask, labels):\n",
    "        return self.model(input_ids, attention_mask=attention_mask, labels=labels)\n",
    "        \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        outputs = self(batch[\"input_ids\"], attention_mask=batch[\"attention_mask\"],\n",
    "                       labels=batch[\"label\"])        \n",
    "        self.log(\"train_loss\", outputs[\"loss\"])\n",
    "        return outputs[\"loss\"]  # this is passed to the optimizer for training\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        outputs = self(batch[\"input_ids\"], attention_mask=batch[\"attention_mask\"],\n",
    "                       labels=batch[\"label\"])        \n",
    "        self.log(\"val_loss\", outputs[\"loss\"], prog_bar=True)\n",
    "        \n",
    "        logits = outputs[\"logits\"]\n",
    "        predicted_labels = torch.argmax(logits, 1)\n",
    "        self.val_acc(predicted_labels, batch[\"label\"])\n",
    "        self.log(\"val_acc\", self.val_acc, prog_bar=True)\n",
    "        \n",
    "    def test_step(self, batch, batch_idx):\n",
    "        outputs = self(batch[\"input_ids\"], attention_mask=batch[\"attention_mask\"],\n",
    "                       labels=batch[\"label\"])        \n",
    "        \n",
    "        logits = outputs[\"logits\"]\n",
    "        predicted_labels = torch.argmax(logits, 1)\n",
    "        self.test_acc(predicted_labels, batch[\"label\"])\n",
    "        self.log(\"accuracy\", self.test_acc, prog_bar=True)\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)\n",
    "        return optimizer\n",
    "    \n",
    "\n",
    "lightning_model = LightningModel(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "2ee664a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from finetuning_scheduler import FinetuningScheduler\n",
    "\n",
    "callbacks = [FinetuningScheduler()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "698b7a7d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/home/jovyan/conda/lib/python3.8/site-packages/finetuning_scheduler/fts_supporters.py:1482: UserWarning: FinetuningScheduler currently depends upon an FTSEarlyStopping callback unless configured in epoch_transitions_only mode. Adding an FTSEarlyStopping callback with default configuration.\n",
      "  rank_zero_warn(\n",
      "/home/jovyan/conda/lib/python3.8/site-packages/finetuning_scheduler/fts_supporters.py:1496: UserWarning: FinetuningScheduler currently depends upon a fine-tuning schedule capable ModelCheckpoint callback such as FTSCheckpoint. Substituting current ModelCheckpoint for FTSCheckpoint\n",
      "  rank_zero_warn(\n",
      "No fine-tuning schedule provided, max_depth set to -1 so iteratively thawing entire model\n",
      "fine-tuning schedule dumped to /home/jovyan/lightning_logs/version_1/LightningModel_ft_schedule.yaml.\n",
      "Generated default fine-tuning schedule '/home/jovyan/lightning_logs/version_1/LightningModel_ft_schedule.yaml' for iterative fine-tuning\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "/home/jovyan/conda/lib/python3.8/site-packages/finetuning_scheduler/fts_supporters.py:1377: UserWarning: FinetuningScheduler configured the provided model to have 2 trainable parameters in phase 0 (the initial training phase) but the optimizer has subsequently been initialized with 102 additional parameters that do not require a gradient. If non-intentional, this state is commonly caused by failing to filter out parameters that do not require a gradient when initializing the optimizer (e.g., `parameters = list(filter(lambda x: x.requires_grad, self.parameters()))` If you intended to initialize the optimizer with parameters that do not require a gradient you may want to ensure they are not included in the 104 parameters that the FinetuningScheduler is currently configured to thaw (sum of all phases) to avoid triggering a  parameter collision and training failure in pytorch during a future fine-tuning phase.\n",
      "  rank_zero_warn(warn_msg)\n",
      "\n",
      "  | Name     | Type                                | Params\n",
      "-----------------------------------------------------------------\n",
      "0 | model    | DistilBertForSequenceClassification | 67.0 M\n",
      "1 | val_acc  | Accuracy                            | 0     \n",
      "2 | test_acc | Accuracy                            | 0     \n",
      "-----------------------------------------------------------------\n",
      "1.5 K     Trainable params\n",
      "67.0 M    Non-trainable params\n",
      "67.0 M    Total params\n",
      "267.820   Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d5097d74e5c74211b4a2456aa4d0f3d7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 0, global step 2917: 'val_loss' reached 0.64435 (best 0.64435), saving model to '/home/jovyan/lightning_logs/version_1/checkpoints/epoch=0-step=2917.ckpt' as top 1\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 1, global step 5834: 'val_loss' reached 0.60712 (best 0.60712), saving model to '/home/jovyan/lightning_logs/version_1/checkpoints/epoch=1-step=5834.ckpt' as top 1\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 2, global step 8751: 'val_loss' reached 0.57646 (best 0.57646), saving model to '/home/jovyan/lightning_logs/version_1/checkpoints/epoch=2-step=8751.ckpt' as top 1\n",
      "`Trainer.fit` stopped: `max_epochs=3` reached.\n"
     ]
    }
   ],
   "source": [
    "trainer = L.Trainer(\n",
    "    max_epochs=3,\n",
    "    callbacks=callbacks,\n",
    "    accelerator=\"gpu\",\n",
    "    devices=[0],\n",
    "    log_every_n_steps=10,\n",
    ")\n",
    "\n",
    "trainer.fit(model=lightning_model,\n",
    "            train_dataloaders=train_loader,\n",
    "            val_dataloaders=val_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "a8eefa9f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "fine-tuning schedule dumped to /home/jovyan/lightning_logs/version_1/LightningModel_ft_schedule.yaml.\n",
      "Restoring states from the checkpoint path at /home/jovyan/lightning_logs/version_1/checkpoints/epoch=2-step=8751.ckpt\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "Loaded model weights from checkpoint at /home/jovyan/lightning_logs/version_1/checkpoints/epoch=2-step=8751.ckpt\n",
      "/home/jovyan/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:488: PossibleUserWarning: Your `test_dataloader`'s sampler has shuffling enabled, it is strongly recommended that you turn shuffling off for val/test/predict dataloaders.\n",
      "  rank_zero_warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a68b6e55608e4e3285adb8cd21ba6c66",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">        Test metric        </span>┃<span style=\"font-weight: bold\">       DataLoader 0        </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">         accuracy          </span>│<span style=\"color: #800080; text-decoration-color: #800080\">    0.7920571565628052     </span>│\n",
       "└───────────────────────────┴───────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m       Test metric       \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m      DataLoader 0       \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[36m \u001b[0m\u001b[36m        accuracy         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m   0.7920571565628052    \u001b[0m\u001b[35m \u001b[0m│\n",
       "└───────────────────────────┴───────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[{'accuracy': 0.7920571565628052}]"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.test(lightning_model, dataloaders=train_loader, ckpt_path=\"best\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "a52d01f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "fine-tuning schedule dumped to /home/jovyan/lightning_logs/version_1/LightningModel_ft_schedule.yaml.\n",
      "Restoring states from the checkpoint path at /home/jovyan/lightning_logs/version_1/checkpoints/epoch=2-step=8751.ckpt\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "Loaded model weights from checkpoint at /home/jovyan/lightning_logs/version_1/checkpoints/epoch=2-step=8751.ckpt\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7e646d0759c54fa0a54cfec976b7f3c6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">        Test metric        </span>┃<span style=\"font-weight: bold\">       DataLoader 0        </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">         accuracy          </span>│<span style=\"color: #800080; text-decoration-color: #800080\">    0.7865999937057495     </span>│\n",
       "└───────────────────────────┴───────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m       Test metric       \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m      DataLoader 0       \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[36m \u001b[0m\u001b[36m        accuracy         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m   0.7865999937057495    \u001b[0m\u001b[35m \u001b[0m│\n",
       "└───────────────────────────┴───────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[{'accuracy': 0.7865999937057495}]"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.test(lightning_model, dataloaders=val_loader, ckpt_path=\"best\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "0106f6f9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "fine-tuning schedule dumped to /home/jovyan/lightning_logs/version_1/LightningModel_ft_schedule.yaml.\n",
      "Restoring states from the checkpoint path at /home/jovyan/lightning_logs/version_1/checkpoints/epoch=2-step=8751.ckpt\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "Loaded model weights from checkpoint at /home/jovyan/lightning_logs/version_1/checkpoints/epoch=2-step=8751.ckpt\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fd91b70174b0412e815143ea9349817d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">        Test metric        </span>┃<span style=\"font-weight: bold\">       DataLoader 0        </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">         accuracy          </span>│<span style=\"color: #800080; text-decoration-color: #800080\">    0.7882999777793884     </span>│\n",
       "└───────────────────────────┴───────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m       Test metric       \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m      DataLoader 0       \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[36m \u001b[0m\u001b[36m        accuracy         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m   0.7882999777793884    \u001b[0m\u001b[35m \u001b[0m│\n",
       "└───────────────────────────┴───────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[{'accuracy': 0.7882999777793884}]"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.test(lightning_model, dataloaders=test_loader, ckpt_path=\"best\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e33b9ea",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
